package main

import (
	"crypto/rand"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"time"
)

// PerformanceResult (已更新)
type PerformanceResult struct {
	Experiment     string
	Parameter      string
	Value          int
	ProofGenTimeMs float64
	VerifyTimeMs   float64
	// TagGenTimeMs        float64
	// AvgProofGenTimeMs   float64
	// TotalProofGenTimeMs float64
	// AvgVerifyTimeMs     float64
	// TotalVerifyTimeMs   float64
}

// 基础功能性实验
func basicExp() {
	log.Println("--- Starting Basic Functionality Experiment (Bilinear Pairing Version) ---")

	log.Println("--- Step 1: System Setup ---")
	const K, FileSize, BlockSize, ChallengedBlocks = 5, 1024 * 16, 1024, 10

	cspIDs := make([]string, K)
	for i := 0; i < K; i++ {
		cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
	}
	params, _ := NewSystemParams()

	log.Println("--- Step 2: Key Generation and Tag Generation ---")
	dataOwner, _ := NewDataOwner(params)
	dummyFile := make([]byte, FileSize)
	rand.Read(dummyFile)

	tagMap, _ := dataOwner.TagGen(dummyFile, "doc1", BlockSize, cspIDs)

	csps := make([]*CSP, K)
	for i, id := range cspIDs {
		csps[i] = NewCSP(id, params, dataOwner)
		csps[i].StoreTags("doc1", tagMap[id])
	}
	log.Printf("All %d CSPs have stored their tags.", K)

	log.Println("\n--- Step 3: Normal Auditing Process ---")
	challenge, _ := dataOwner.GenChallenge("doc1", ChallengedBlocks)

	partialProofs := make([]*PartialProof, K)
	for i, csp := range csps {
		proof, _ := csp.PartialProofGen(challenge)
		partialProofs[i] = proof
	}

	verifierCSP := csps[0]
	finalProof := verifierCSP.AggregateProof(partialProofs)
	isVerifiedNormal := verifierCSP.VerifyProof(finalProof, challenge, dataOwner.pk, cspIDs)

	if isVerifiedNormal {
		log.Println("✅SUCCESS: Verification passed in the normal process.")
	} else {
		log.Println("❌FAILURE: Verification failed in the normal process.")
	}

	log.Println("\n--- Step 4: Auditing Process After Data Tampering ---")
	tamperedBlockIndex := challenge.Indices[0]
	log.Printf("!!! Tampering data block %d on ALL CSPs (simulated) !!!", tamperedBlockIndex)
	tamperedValue := params.Pairing.NewZr().Set1()
	originalBlock := params.Pairing.NewZr().Set(dataOwner.originalData["doc1"][tamperedBlockIndex])
	dataOwner.originalData["doc1"][tamperedBlockIndex].Add(originalBlock, tamperedValue)

	tamperedPartialProofs := make([]*PartialProof, K)
	for i, csp := range csps {
		proof, _ := csp.PartialProofGen(challenge)
		tamperedPartialProofs[i] = proof
	}
	tamperedFinalProof := verifierCSP.AggregateProof(tamperedPartialProofs)
	isTamperingDetected := !verifierCSP.VerifyProof(tamperedFinalProof, challenge, dataOwner.pk, cspIDs)

	if isTamperingDetected {
		log.Println("✅SUCCESS: Data tampering was successfully detected.")
	} else {
		log.Println("❌FAILURE: Data tampering was NOT detected.")
	}
	dataOwner.originalData["doc1"][tamperedBlockIndex].Set(originalBlock)
}

// 实验1：考察文件大小的影响
func fileSizeExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Running Experiment 1: Performance vs. File Size >>>>>>>>>>")
	const K = 5
	const NumBlocks = 4096
	fileSizeOptionsMB := []int{4, 8, 12, 16, 20}

	cspIDs := make([]string, K)
	for i := range K {
		cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
	}
	params, _ := NewSystemParams()
	dataOwner, _ := NewDataOwner(params)

	for _, sizeMB := range fileSizeOptionsMB {
		fileSize := sizeMB * 1024 * 1024
		blockSize := fileSize / NumBlocks
		log.Printf("\n--- Testing with FileSize: %dMB (NumBlocks: %d, BlockSize: %d bytes) ---", sizeMB, NumBlocks, blockSize)

		dummyFile := make([]byte, fileSize)
		rand.Read(dummyFile)

		tagMap, _ := dataOwner.TagGen(dummyFile, "perf_doc", blockSize, cspIDs)

		csps := make([]*CSP, K)
		for i, id := range cspIDs {
			csps[i] = NewCSP(id, params, dataOwner)
			csps[i].StoreTags("perf_doc", tagMap[id])
		}

		challenge, _ := dataOwner.GenChallenge("perf_doc", NumBlocks)

		// TODO: 生成部分证明 和 完整证明的时间都算上
		start := time.Now()
		partialProofs := make([]*PartialProof, K)
		for i, csp := range csps {
			proof, _ := csp.PartialProofGen(challenge)
			partialProofs[i] = proof
		}
		partialProofGenTime := time.Since(start) / time.Duration(K)
		start = time.Now()
		finalProof := csps[0].AggregateProof(partialProofs)
		proofGenTime := time.Since(start) + partialProofGenTime

		// TODO: 验证一次
		startVerify := time.Now()
		csps[0].VerifyProof(finalProof, challenge, dataOwner.pk, cspIDs)
		verifyTime := time.Since(startVerify)

		result := PerformanceResult{
			Experiment: "FileSize", Parameter: "FileSize(MB)", Value: sizeMB,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
		})
		writer.Flush()
	}
}

// 实验2：考察 CSP 数量的影响
func cspCountExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Running Experiment 2: Performance vs. Number of CSPs >>>>>>>>>>")
	const NumBlocks = 4096
	const BlockSize = 1024
	const FileSize = NumBlocks * BlockSize
	cspCounts := []int{5, 10, 15, 20, 25, 30}

	dummyFile := make([]byte, FileSize)
	rand.Read(dummyFile)
	params, _ := NewSystemParams()
	dataOwner, _ := NewDataOwner(params)

	for _, k := range cspCounts {
		log.Printf("\n--- Testing with %d CSPs ---", k)

		cspIDs := make([]string, k)
		for i := 0; i < k; i++ {
			cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
		}

		tagMap, _ := dataOwner.TagGen(dummyFile, "perf_doc", BlockSize, cspIDs)

		csps := make([]*CSP, k)
		for i, id := range cspIDs {
			csps[i] = NewCSP(id, params, dataOwner)
			csps[i].StoreTags("perf_doc", tagMap[id])
		}

		challenge, _ := dataOwner.GenChallenge("perf_doc", NumBlocks)

		// TODO: 生成部分证明 和 完整证明的时间都算上
		start := time.Now()
		partialProofs := make([]*PartialProof, k)
		for i, csp := range csps {
			proof, _ := csp.PartialProofGen(challenge)
			partialProofs[i] = proof
		}
		partialProofGenTime := time.Since(start) / time.Duration(k)
		start = time.Now()
		finalProof := csps[0].AggregateProof(partialProofs)
		proofGenTime := time.Since(start) + partialProofGenTime

		// TODO: 验证一次
		startVerify := time.Now()
		csps[0].VerifyProof(finalProof, challenge, dataOwner.pk, cspIDs)
		verifyTime := time.Since(startVerify)

		result := PerformanceResult{
			Experiment: "CSPCount", Parameter: "NumCSPs", Value: k,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
		})
		writer.Flush()
	}
}

// 实验3：考察文件分块数量的影响
func numBlocksExp(writer *csv.Writer) {
	log.Println("\n\n<<<<<<<<<< Running Experiment 3: Performance vs. Number of Blocks >>>>>>>>>>")
	const K = 5
	const FileSize = 4 * 1024 * 1024
	numBlocksOptions := []int{1024, 2048, 4096, 8192, 16384}

	cspIDs := make([]string, K)
	for i := 0; i < K; i++ {
		cspIDs[i] = fmt.Sprintf("CSP-%d", i+1)
	}
	params, _ := NewSystemParams()
	dataOwner, _ := NewDataOwner(params)

	for _, numBlocks := range numBlocksOptions {
		blockSize := FileSize / numBlocks
		log.Printf("\n--- Testing with %d blocks (FileSize: %dMB, BlockSize: %d bytes) ---", numBlocks, FileSize/(1024*1024), blockSize)

		dummyFile := make([]byte, FileSize)
		rand.Read(dummyFile)

		tagMap, _ := dataOwner.TagGen(dummyFile, "perf_doc", blockSize, cspIDs)

		csps := make([]*CSP, K)
		for i, id := range cspIDs {
			csps[i] = NewCSP(id, params, dataOwner)
			csps[i].StoreTags("perf_doc", tagMap[id])
		}

		challenge, _ := dataOwner.GenChallenge("perf_doc", numBlocks)

		// TODO: 生成部分证明 和 完整证明的时间都算上
		start := time.Now()
		partialProofs := make([]*PartialProof, K)
		for i, csp := range csps {
			proof, _ := csp.PartialProofGen(challenge)
			partialProofs[i] = proof
		}
		partialProofGenTime := time.Since(start) / time.Duration(K)
		start = time.Now()
		finalProof := csps[0].AggregateProof(partialProofs)
		proofGenTime := time.Since(start) + partialProofGenTime

		// TODO: 验证一次
		startVerify := time.Now()
		csps[0].VerifyProof(finalProof, challenge, dataOwner.pk, cspIDs)
		verifyTime := time.Since(startVerify)

		result := PerformanceResult{
			Experiment: "NumBlocks", Parameter: "NumberOfBlocks", Value: numBlocks,
			ProofGenTimeMs: float64(proofGenTime.Microseconds()) / 1000.0,
			VerifyTimeMs:   float64(verifyTime.Microseconds()) / 1000.0,
		}
		log.Printf("[PERF] Result: %+v", result)
		writer.Write([]string{
			result.Experiment, result.Parameter, fmt.Sprintf("%d", result.Value),
			fmt.Sprintf("%.3f", result.ProofGenTimeMs),
			fmt.Sprintf("%.3f", result.VerifyTimeMs),
		})
		writer.Flush()
	}
}

func main() {
	logFile, err := os.OpenFile("experiment_exp3-pbc.log", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	if err != nil {
		log.Fatalf("Failed to open log file: %v", err)
	}
	defer logFile.Close()
	mw := io.MultiWriter(os.Stdout, logFile)
	log.SetOutput(mw)

	csvFile, err := os.Create("results_exp3-pbc.csv")
	if err != nil {
		log.Fatalf("Failed to create CSV file: %v", err)
	}
	defer csvFile.Close()
	csvWriter := csv.NewWriter(csvFile)
	defer csvWriter.Flush()
	csvWriter.Write([]string{"Experiment", "Parameter","Value","ProofGenTime(ms)", "VerifyTime(ms)"})

	// basicExp()
	fileSizeExp(csvWriter)
	cspCountExp(csvWriter)
	numBlocksExp(csvWriter)

	log.Println("\n\n<<<<<<<<<< All experiments for exp3-pbc finished. >>>>>>>>>>")
}
